from keras.models import Sequential
from keras.layers import Convolution2D, ZeroPadding2D, MaxPooling2D
from keras.layers.core import Flatten, Dense, Dropout, Lambda
from keras import backend as K
import h5py
from keras.optimizers import SGD

def global_average_pooling(x):
    # return K.max(x, axis=(2,3)) 
 
    # normalize_rate = 0.9
    # epsilon = 1e-10
    # max_val = K.max(x, axis=(2,3), keepdims = True)
    # normalized = x / (normalize_rate * max_val + epsilon)
    # square_normalized = K.square(normalized)
    # return K.mean(square_normalized, axis = (2,3))

    normalize_rate = 1
    epsilon = 1e-10
    sum_val = K.sum(x, axis=(2,3), keepdims = True)
    normalized = x / (sum_val * normalize_rate + epsilon)
    square_normalized = x * normalized
    return K.mean(square_normalized, axis = (2,3))

    # normalize_rate = 1
    # epsilon = 1e-10
    # square = K.square(x)
    # sum_val = K.sum(square, axis=(2,3), keepdims = True)
    # normalized = x / (sum_val * normalize_rate + epsilon)
    # square_normalized = x * normalized
    # return K.mean(square_normalized, axis = (2,3))

    # return K.mean(x, axis = (2, 3))


def global_average_pooling_shape(input_shape):
    return input_shape[0:2]

def VGG16_convolutions():
    #K.set_image_data_format('channels_first')
    model = Sequential()
    model.add(ZeroPadding2D((1,1),input_shape=(3,None,None)))
    model.add(Convolution2D(64, 3, 3, activation='relu', name='conv1_1'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(64, 3, 3, activation='relu', name='conv1_2'))
    model.add(MaxPooling2D((2, 2), strides=(2, 2)))

    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(128, 3, 3, activation='relu', name='conv2_1'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(128, 3, 3, activation='relu', name='conv2_2'))
    model.add(MaxPooling2D((2, 2), strides=(2, 2)))

    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_1'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_2'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(256, 3, 3, activation='relu', name='conv3_3'))
    model.add(MaxPooling2D((2, 2), strides=(2, 2)))

    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_1'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_2'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv4_3'))
    model.add(MaxPooling2D((2, 2), strides=(2, 2)))

    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv5_1'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv5_2'))
    model.add(ZeroPadding2D((1, 1)))
    model.add(Convolution2D(512, 3, 3, activation='relu', name='conv5_3'))
    return model

def get_model():
    model = VGG16_convolutions()

    model = load_model_weights(model, "vgg16_weights.h5")
    
    # for layer in model.layers[:-4]:
    #     layer.trainable = False
    for layer in model.layers[-4:]:
        layer.trainable = True
    
    model.add(Lambda(global_average_pooling, 
              output_shape=global_average_pooling_shape))
    model.add(Dense(2, activation = 'softmax', init='uniform'))
    sgd = SGD(lr=0.001, decay=1e-6, momentum=0.5, nesterov=True)
    model.compile(loss = 'categorical_crossentropy', optimizer = sgd, metrics=['accuracy'])
    print(model.summary())
    return model

def load_model_weights(model, weights_path):
    print ('Loading model.')
    f = h5py.File(weights_path)
    for k in range(f.attrs['nb_layers']):
        if k >= len(model.layers):
            # we don't look at the last (fully-connected) layers in the savefile
            break
        g = f['layer_{}'.format(k)]
        weights = [g['param_{}'.format(p)] for p in range(g.attrs['nb_params'])]
        model.layers[k].set_weights(weights)
        model.layers[k].trainable = False
    f.close()
    print ('Model loaded.')
    return model

def get_output_layer(model, layer_name):
    # get the symbolic outputs of each "key" layer (we gave them unique names).
    layer_dict = dict([(layer.name, layer) for layer in model.layers])
    layer = layer_dict[layer_name]
    return layer